import os.path
import time
import traceback

from Lib.zcm.assistant import FwVariables
from robot.api import logger

try:
    import paramiko
except:
    pass
import os


class SFTPOperate:
    def __init__(self):
        pass

    def ftp_connect(self, link_name, **kwargs):
        sftp_host = FwVariables.get_sftp_host(link_name)
        username = FwVariables.get_sftp_user(link_name)
        password = FwVariables.get_sftp_password(link_name)
        port = FwVariables.get_sftp_port(link_name)
        self.sf = paramiko.Transport((sftp_host, port))
        self.sf.connect(username=username, password=password, **kwargs)
        self.sftp = paramiko.SFTPClient.from_transport(self.sf)

    def sftp_upload(self, link_name, local_path, remote_path, **kwargs):
        try:
            local_path = local_path.replace(os.sep, "/")
            self.ftp_connect(link_name, **kwargs)
            home_dir = FwVariables.get_sftp_home_dir(link_name)
            if os.path.isdir(local_path):  # 判断本地参数是目录还是文件
                local_base_path = os.path.basename(local_path)
                for f in os.listdir(local_path):  # 遍历本地目录
                    child_local_path = os.path.join(local_path, f)
                    if os.path.isdir(child_local_path):
                        mk_dir_path = os.path.join(home_dir, os.path.join(remote_path, local_base_path)).replace(os.sep,
                                                                                                                 "/")
                        try:
                            self.sftp.stat(mk_dir_path)
                        except:
                            self.sftp.mkdir(mk_dir_path)
                        self.sftp_upload(link_name, child_local_path, mk_dir_path)
                    else:
                        mk_dir_path = os.path.join(home_dir,
                                                   os.path.join(remote_path, local_base_path).replace(os.sep, "/"))
                        try:
                            self.sftp.stat(mk_dir_path)
                        except:
                            self.sftp.mkdir(mk_dir_path)
                        remote = os.path.join(home_dir, os.path.join(remote_path, local_base_path, f))
                        remote = remote.replace(os.sep, "/")
                        self.sftp.put(child_local_path, remote)  # 上传目录中的文件
            else:
                remote = os.path.join(home_dir, remote_path).replace(os.sep, "/")
                dir_path = os.path.dirname(remote)
                try:
                    self.sftp.stat(dir_path)
                except:
                    self.sftp.mkdir(dir_path)
                self.sftp.put(local_path, remote)  # 上传文件
        except Exception as e:
            # logger.error(traceback.format_exc())
            raise Exception(e)
        finally:
            self.sf.close()

    def sftp_download(self, link_name, local_path, remote_path, **kwargs):
        try:
            self.ftp_connect(link_name, **kwargs)
            home_dir = FwVariables.get_sftp_home_dir(link_name)
            remote_path = os.path.join(home_dir, remote_path).replace(os.sep, "/")
            # if os.path.isdir(local_path):  # 判断本地参数是目录还是文件
            #     if self.sftp.listdir()
            #     for f in self.sftp.listdir(remote_path):  # 遍历远程目录
            #         sftp.get(os.path.join(remote + f), os.path.join(local + f))  # 下载目录中文件
            # else:
            self.sftp.get(remote_path, local_path)  # 下载文件
        except Exception:
            logger.error(traceback.format_exc())
        finally:
            self.sf.close()


def send_command(shell, command, wait_time=1, time_out=5):
    """
    发送命令并接收输出。
    """
    shell.send(command + '\n')
    time.sleep(wait_time)  # 等待命令执行
    flag = 0
    while True:
        if flag > time_out:
            break
        if shell.recv_ready():
            result = shell.recv(10000).decode('utf-8')
            logger.info(result)
            return result
        else:
            time.sleep(1)
            flag += 1


def execute_remote_command(link_name, command_info, wait_time):
    bastion_host = FwVariables.get_remote_host(link_name)
    bastion_user = FwVariables.get_remote_user(link_name)
    bastion_password = FwVariables.get_remote_password(link_name)
    bastion_port = FwVariables.get_remote_port(link_name)
    time_out = FwVariables.get_remote_timeout(link_name)

    # 创建SSH客户端
    ssh = paramiko.SSHClient()
    ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
    result = ""
    try:
        # 连接到堡垒机
        ssh.connect(bastion_host, username=bastion_user, password=bastion_password, port=bastion_port)
        shell = ssh.invoke_shell()
        shell.settimeout(time_out)
        for command in command_info:
            time.sleep(int(wait_time))
            inputs = command_info[command]
            # 等待命令提示符
            result += send_command(shell, command, time_out)
            if not inputs:
                continue
            for i in inputs:
                result += send_command(shell, i, time_out)
    finally:
        ssh.close()  # 关闭连接
    return result


sftp_obj = SFTPOperate()

if __name__ == '__main__':
    # sftp_obj = SFTPOperate()
    # sftp_obj.sftp_upload("ftp", "D:\\ftp\\1.xls", "/home/ocs90/ftp/2.xls")
    # sftp_obj.sftp_download("ftp", "D:\\ftp\\xxxxxxxxxxx.xls", "/home/ocs90/ftp/2.xls")

    # a = execute_remote_command("ZMQ", '/var/lib/zmq/namesrv/zmq/bin/mqadmin producer -t testlfm', ["12313123131", "22222222"], need_manual_terminate=True)
    print(execute_remote_command("ZMQ", 'ps -ef', []))
